{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8d70e8fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import numpy as np\n",
    "import pynq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d17a0bc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "bitfile = \"conv.bit\"\n",
    "overlay = pynq.Overlay(bitfile)\n",
    "# !./dfs 0 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "82683eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "dma = overlay.axi_dma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e9f02b07",
   "metadata": {},
   "outputs": [],
   "source": [
    "height, width, in_channels = 5, 5, 8\n",
    "out_channels = 1\n",
    "P_ICH, P_OCH = 4, 1\n",
    "kernel_size, stride, padding = 3, 1, 1\n",
    "A_BIT, W_BIT, B_BIT = 8, 8, 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c6c43419",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_size = height * width * in_channels\n",
    "input_data = (np.arange(input_size) % 128).astype(np.uint8)\n",
    "input_data = input_data.reshape(height, width, in_channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "0cc5148d",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_buffer = pynq.allocate(shape=(in_channels * height * width * kernel_size * kernel_size), dtype=np.uint8, cacheable=1)\n",
    "output_buffer = pynq.allocate(shape=(out_channels * height * width), dtype=np.uint32, cacheable=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a3c88668",
   "metadata": {},
   "outputs": [],
   "source": [
    "def im2col(input_data, kernel_size, stride=1, padding=0):\n",
    "    input_height, input_width, channels = input_data.shape\n",
    "    out_height = (input_height + 2 * padding - kernel_size) // stride + 1\n",
    "    out_width  = (input_width  + 2 * padding - kernel_size) // stride + 1\n",
    "    output_data = np.zeros((out_height * out_width, kernel_size * kernel_size * channels), dtype=input_data.dtype)\n",
    "\n",
    "    for oh in range(out_height):\n",
    "        for ow in range(out_width):\n",
    "            for c in range(channels):\n",
    "                for kh in range(kernel_size):\n",
    "                    for kw in range(kernel_size):\n",
    "                        ih = oh * stride - padding + kh\n",
    "                        iw = ow * stride - padding + kw\n",
    "                        if 0 <= ih < input_height and 0 <= iw < input_width:\n",
    "                            output_data[oh * out_width + ow, kh * kernel_size * channels + kw * channels + c] = input_data[ih, iw, c]\n",
    "                        else:\n",
    "                            output_data[oh * out_width + ow, kh * kernel_size * channels + kw * channels + c] = 0\n",
    "    return output_data.reshape(out_height * out_width, -1)\n",
    "\n",
    "def reshape_im2col_for_hardware(im2col_data, kernel_size, in_channels, P_ICH):\n",
    "    im2col_reshaped = im2col_data.reshape(-1, kernel_size, kernel_size, in_channels // P_ICH, P_ICH)\n",
    "    im2col_reshaped = im2col_reshaped.transpose(0, 3, 1, 2, 4)\n",
    "    return im2col_reshaped.reshape(-1, kernel_size * kernel_size * in_channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "c6776a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "im2col_data = im2col(input_data, kernel_size, stride, padding)\n",
    "im2col_reshaped = reshape_im2col_for_hardware(im2col_data, kernel_size, in_channels, P_ICH)\n",
    "im2col_reshaped = im2col_reshaped.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "af2634e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.copyto(input_buffer, im2col_reshaped)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "392f5481",
   "metadata": {},
   "outputs": [],
   "source": [
    "dma.sendchannel.transfer(input_buffer)\n",
    "dma.recvchannel.transfer(output_buffer)\n",
    "dma.recvchannel.wait()\n",
    "dma.sendchannel.wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "fd0ec34b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PynqBuffer([ 53680,  85640, 103880, 122120,  80048, 121800, 181164,\n",
       "            201612, 222060, 139848, 128520, 153356, 121068, 141516,\n",
       "             86664,  98888, 121964, 114252, 134700,  82632,  37296,\n",
       "             43912,  49352,  58376,  35504], dtype=uint32)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3aae6f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
